package com.xuzimian.ai.ocr.paddle

import ai.djl.inference.Predictor
import ai.djl.modality.Classifications
import ai.djl.modality.Classifications.Classification
import ai.djl.modality.cv.Image
import ai.djl.modality.cv.ImageFactory
import ai.djl.modality.cv.output.BoundingBox
import ai.djl.modality.cv.output.DetectedObjects
import ai.djl.modality.cv.output.DetectedObjects.DetectedObject
import ai.djl.modality.cv.output.Rectangle
import ai.djl.modality.cv.util.NDImageUtils
import ai.djl.ndarray.NDManager
import ai.djl.paddlepaddle.zoo.cv.imageclassification.PpWordRotateTranslator
import ai.djl.paddlepaddle.zoo.cv.objectdetection.PpWordDetectionTranslator
import ai.djl.paddlepaddle.zoo.cv.wordrecognition.PpWordRecognitionTranslator
import ai.djl.repository.zoo.Criteria
import java.nio.file.Files
import java.nio.file.Paths
import java.util.concurrent.ConcurrentHashMap
import kotlin.io.path.Path


class ApiApplication

private val classLoader = ApiApplication::class.java.classLoader


fun main(args: Array<String>) {
    val img = getImage("https://resources.djl.ai/images/flight_ticket.jpg")
//    img.wrappedImage

    // 文字识别 最终结果
    val recognizer = getTextRecognizer()

    // 获取文字区域检测
    val detector = getTextAreaDetector()
    val detectedObj = detector.predict(img)
    println(detectedObj);

    // 尝试画出文字区域框，并输出
    val newImage = img.duplicate()
    // img.duplicate(Image.Type.TYPE_INT_ARGB);
    newImage.drawBoundingBoxes(detectedObj)
    //  newImage.wrappedImage
    newImage.save(Files.newOutputStream(Paths.get("build").resolve("detected-1.png")), "png")


    // 文字角度识别器
    val rotateClassifier = getTextRotateClassifier()
    val result: Classification = rotateClassifier.predict(img).best()

    // 尝试提取一个文字块
    val boxes: List<DetectedObject> = detectedObj.items()
    val sample = getSubImage(img, boxes[5].getBoundingBox())
    // 保存识别的小块区域
    sample.save(Files.newOutputStream(Paths.get("build").resolve("detected5-2.png")), "png")
    // 输出 识别小块区域
    println(recognizer.predict(sample))


    val taggedImage = markTextOnImage(img, rotateClassifier, recognizer, boxes)
    taggedImage.save(Files.newOutputStream(Paths.get("build").resolve("detected.png")), "png")
}

/**
 * 最后，让我们在整个图像上运行这些模型并查看结果。 DJL 提供了丰富的图像工具包，允许您在图像上绘制文本并显示它们。
 */
private fun markTextOnImage(
    img: Image,
    rotateClassifier: Predictor<Image, Classifications>,
    recognizer: Predictor<Image, String>,
    boxes: List<DetectedObject>
): Image {
    val names: MutableList<String> = ArrayList()
    val prob: MutableList<Double> = ArrayList()
    val rect: MutableList<BoundingBox> = ArrayList()
    val newImage = img.duplicate()

    for (i in boxes.indices) {
        var subImg = getSubImage(img, boxes[i].boundingBox)
        if (subImg.height * 1.0 / subImg.width > 1.5) {
            subImg = rotateImg(subImg)
        }
        val result: Classification = rotateClassifier.predict(subImg).best()
        if ("Rotate" == result.className && result.probability > 0.8) {
            subImg = rotateImg(subImg)
        }
        val name: String = recognizer.predict(subImg)
        println("已识别文字：$name")
        names.add(name)
        prob.add(-1.0)
        rect.add(boxes[i].boundingBox)
    }
    newImage.drawBoundingBoxes(DetectedObjects(names, prob, rect))
    return newImage
}


/**
 * 文字区域检测: 单词检测模型中，我们加载从 PaddleOCR 导出的模型。 之后，我们可以从中生成一个名为检测器的 DJL 预测器。
 *  - 可以从中检测出单词块。 模型的原始输出是标记所有单词区域的位图。 PpWordDetectionTranslator 将输出位图转换为矩形边界框，以便我们裁剪图像。
 *  - 可以看到的，单词块非常狭窄，不包括所有单词的整体。 让我们尝试稍微扩展一下以获得更好的结果。 extendRect 将框的高度和宽度扩展至一定比例。 getSubImage 将裁剪图像并提取单词块。
 *  - model: https://resources.djl.ai/test-models/paddleOCR/mobile/det_db.zip
 */
fun getTextAreaDetector(): Predictor<Image, DetectedObjects> {
    val modelPath = classLoader.getResource("det_db.zip")!!.path
    var criteria1 = Criteria.builder()
        .optEngine("PaddlePaddle")
        .setTypes(Image::class.java, DetectedObjects::class.java)
        .optModelPath(Path(modelPath))
        .optTranslator(PpWordDetectionTranslator(ConcurrentHashMap<String, String>()))
        .build();

    var detectionModel = criteria1.loadModel();
    return detectionModel.newPredictor();
}


/**
 * 文字角度检测
 * 该模型是从 PaddleOCR 导出的，可以帮助识别图像是否需要旋转。 以下代码将加载此模型并创建一个 rotateClassifier。
 * - model: https://resources.djl.ai/test-models/paddleOCR/mobile/cls.zip
 */
fun getTextRotateClassifier(): Predictor<Image, Classifications> {
    val modelPath = classLoader.getResource("cls.zip")!!.path
    var criteria2 = Criteria.builder()
        .optEngine("PaddlePaddle")
        .setTypes(Image::class.java, Classifications::class.java)
        .optModelPath(Path(modelPath))
        .optTranslator(PpWordRotateTranslator())
        .build();
    var rotateModel = criteria2.loadModel();
    return rotateModel.newPredictor();
}


/**
 * 文字识别
 * 文字识别模型是从PaddleOCR导出的，可以识别图像上的文字。
 * https://resources.djl.ai/test-models/paddleOCR/mobile/rec_crnn.zip
 */
fun getTextRecognizer(): Predictor<Image, String> {
    val modelPath = classLoader.getResource("rec_crnn.zip")!!.path
    var criteria3 = Criteria.builder()
        .optEngine("PaddlePaddle")
        .setTypes(Image::class.java, String::class.java)
        .optModelPath(Path(modelPath))
        .optTranslator(PpWordRecognitionTranslator())
        .build();
    var recognitionModel = criteria3.loadModel()
    return recognitionModel.newPredictor()
}


/**
 *
 * 将裁剪图像并提取单词块。
 */
fun getSubImage(img: Image, box: BoundingBox): Image {
    val rect: Rectangle = box.getBounds()
    val extended: DoubleArray = extendRect(rect.getX(), rect.getY(), rect.getWidth(), rect.getHeight())
    val width = img.width
    val height = img.height
    val recovered = intArrayOf(
        (extended[0] * width).toInt(),
        (extended[1] * height).toInt(),
        (extended[2] * width).toInt(),
        (extended[3] * height).toInt()
    )

    return img.getSubImage(recovered[0], recovered[1], recovered[2], recovered[3])
}


/**
 * 扩展文字框的大小,将框的高度和宽度扩展至一定比例。
 */
fun extendRect(xmin: Double, ymin: Double, width: Double, height: Double): DoubleArray {
    var width = width
    var height = height
    val centerx = xmin + width / 2
    val centery = ymin + height / 2
    if (width > height) {
        width += height * 2.0
        height *= 3.0
    } else {
        height += width * 2.0
        width *= 3.0
    }
    val newX: Double = if (centerx - width / 2 < 0) 0.0 else centerx - width / 2
    val newY: Double = if (centery - height / 2 < 0) 0.0 else centery - height / 2
    val newWidth = if (newX + width > 1) 1 - newX else width
    val newHeight = if (newY + height > 1) 1 - newY else height
    return doubleArrayOf(newX, newY, newWidth, newHeight)
}

/*
  旋转图像
 */
fun rotateImg(image: Image): Image {
    NDManager.newBaseManager().use { manager ->
        val rotated = NDImageUtils.rotate90(image.toNDArray(manager), 1)
        return ImageFactory.getInstance().fromNDArray(rotated)
    }
}

private fun getImage(url: String): Image {
    return ImageFactory.getInstance().fromUrl(url)
}



